from functionSIm_dataPre import functionSimDataset
from torch.utils.data import DataLoader
from functionSimModel import functionSim
import torch
from tqdm import tqdm
from datetime import datetime
from sklearn.metrics import roc_auc_score
from torch import nn
from functionSim_config import *
import shelve


def write_to_logs(inf):
    with open("/home/cyw/projects/function_sim_project/all_logs/EECG_train_logs.txt", 'a+') as file:
        data = str(datetime.now())+"\t"+inf+"\n"
        file.write(data)
        print(data)


class functionModelWarpper():
    def __init__(self, model=None) -> None:
        # seed_everything()
        if model == None:
            self.model = functionSim(
                embSize=embSize, wordSize=wordsize, layerSize=laySize, depth=depth)
            write_to_logs("functionSim      use_hete:{}    use_cross:{}\n".format(
                use_heterogeous, str(use_cross_interaction)+" "+cross_interaction_name))
            write_to_logs("模型参数：lr:{}    epoch:{}    laysize:{}  depth:{}    embsize:{}    margin:{}\n".format(
                lr, epoch, laySize, depth, embSize, margin))
        else:
            self.model = model

        opt = [torch.optim.Adam([{'params': self.model.parameters()},], lr=lr),
               torch.optim.SGD(self.model.parameters(), lr=lr, momentum=0.9),
               torch.optim.RMSprop(self.model.parameters(), lr=lr),
               torch.optim.Adagrad(self.model.parameters(), lr=lr),
               torch.optim.Adadelta(self.model.parameters(), lr=1.0)]

        my_loss = [nn.MSELoss(),
                   nn.CrossEntropyLoss()
                   ]
        # Triplet Loss 和 Contrastive Loss 一般在训练过程中计算

        self.optimizer = opt[0]

    def get_model_dataloader(self, datas, name, is_train, is_shuffle=True):
        """
            else_datas用于加载其它的输入,方便其他模型的对比
        """
        my_dataset = functionSimDataset(datas, name, is_train=is_train)
        dataloader = DataLoader(my_dataset, batch_size=batchSize, shuffle=is_shuffle, num_workers=16,
                                collate_fn=my_dataset.adjust_samples_to_same_dimension)
        return dataloader

    def train(self, dataloader, valid_dataloader, validValues):
        print("----------------开始训练 use hete:{} use cross:{}----------------".format(
            use_heterogeous, use_cross_interaction))
        totalLoss = "-"
        best_auc = -float("inf")
        for e in tqdm(range(epoch)):
            self.model.train()
            temploss = 0
            for i, (query, target1, target2, target3) in enumerate(dataloader):
                torch.cuda.empty_cache()
                self.optimizer.zero_grad()

                adj_x, att_x, vtype_x = target1
                adj_y, att_y, vtype_y = target2
                adj_z, att_z, vtype_z = target3

                adj_x, adj_y, adj_z = adj_x.to(
                    device), adj_y.to(device), adj_z.to(device)
                att_x, att_y, att_z = att_x.to(
                    device), att_y.to(device), att_z.to(device)
                vtype_x, vtype_y, vtype_z = vtype_x.to(
                    device), vtype_y.to(device), vtype_z.to(device)

                score1 = self.model(adj_x, att_x, vtype_x,
                                    adj_y, att_y, vtype_y)  # B*1
                score2 = self.model(adj_x, att_x, vtype_x,
                                    adj_z, att_z, vtype_z)

                #   triplet loss
                loss = torch.relu(score2 - score1 + margin)
                loss = torch.mean(loss)

                temploss += loss
                loss.backward()
                self.optimizer.step()
            temp_auc = self.test_model_auc(valid_dataloader, validValues)
            if show_test_auc:
                temp_auc1 = self.test_model_auc(test_dataloader, testvalues)
                write_to_logs("epoch:{} auc:{} test auc:{} loss:{}".format(
                    e, temp_auc, temp_auc1, temploss))
            else:
                write_to_logs("epoch:{} auc:{} loss:{}".format(
                    e, temp_auc, temploss))
            if temp_auc > best_auc:
                best_auc = temp_auc
                torch.save(
                    self.model, r'/home/cyw/projects/function_sim_project/all_data/models/EECG_model.pth')
                write_to_logs("模型已保存")
            totalLoss += str(float(temploss))
            totalLoss += "-"
        print("训练完成")

    def test_model_auc(self, dataloader, values):
        ans = self.get_model_score(dataloader)
        auc_score = roc_auc_score(values, ans)
        return auc_score

    def get_model_score(self, dataloader):
        """
            输入指定格式数据，通过训练好的模型得到相应的得分
            把dataloader输入进来，方便之前做数据分布的处理
            不为train了，则数据只有前三个是有用的，第四个出现值说明代码错误
        """
        self.model.eval()
        with torch.no_grad():
            res = []
            for i, (query, target1, target2, target3) in enumerate(dataloader):
                torch.cuda.empty_cache()
                self.optimizer.zero_grad()
                adj_x, att_x, vtype_x = target1
                adj_y, att_y, vtype_y = target2
                adj_x = adj_x.to(device)
                adj_y = adj_y.to(device)
                att_x = att_x.to(device)
                att_y = att_y.to(device)
                vtype_x = vtype_x.to(device)
                vtype_y = vtype_y.to(device)
                score1 = self.model(adj_x, att_x, vtype_x,
                                    adj_y, att_y, vtype_y)  # B*1
                res.append(score1)
            ans = torch.cat((res))
        return ans.tolist()

    def get_pair_infs(self, with_Lable=False, selectPath="False"):
        """
            获取之前生成的的训练对信息
            with_Lable
                True：返回(sample1,sample2,flag,sample1_famlable,sample2_famlable)的数组
                False:返回(sample1,sample2,flag)的数组
            selectPath，用于样本对的存放地址
        """
        pairInfPath = r"/home/cyw/projects/function_sim_project/all_data/EECG_pair_infs/pair_infs"
        if selectPath == "False":
            with shelve.open(pairInfPath) as file:
                trainPair = file["trainPair"]
                testPair = file["testPair"]
                validPair = file["validPair"]
        else:
            with shelve.open(selectPath) as file:
                trainPair = file["trainPair"]
                testPair = file["testPair"]
                validPair = file["validPair"]
        res = {}
        res["train"], res["test"], res["valid"] = [], [], []
        for i in range(len(trainPair)):
            temp = trainPair[i][:3] if with_Lable == False else trainPair
            res["train"].append(temp)
        for i in range(len(testPair)):
            temp = testPair[i][:3] if with_Lable == False else testPair
            res["test"].append(temp)
        for i in range(len(validPair)):
            temp = validPair[i][:3] if with_Lable == False else validPair
            res["valid"].append(temp)
        return res


if __name__ == "__main__":
    funcTrainer = functionModelWarpper()
    data_pair_infs = funcTrainer.get_pair_infs()

    testvalues, validvalues = [], []
    for i in range(len(data_pair_infs["test"])):
        testvalues.append(data_pair_infs["test"][i][2])
    for i in range(len(data_pair_infs["valid"])):
        validvalues.append(data_pair_infs["valid"][i][2])

    train_dataloader = funcTrainer.get_model_dataloader(
        data_pair_infs["train"], "train", is_train=True)
    valid_dataloader = funcTrainer.get_model_dataloader(
        data_pair_infs["valid"], "valid", is_train=False, is_shuffle=False)
    test_dataloader = funcTrainer.get_model_dataloader(
        data_pair_infs["test"], "test", is_train=False, is_shuffle=False)
    funcTrainer.train(train_dataloader, valid_dataloader, validvalues)
